import paddle_aux
import paddle
"""
Loss functions
"""
from utils.metrics import bbox_iou
from utils.torch_utils import is_parallel


def smooth_BCE(eps=0.1):
    return 1.0 - 0.5 * eps, 0.5 * eps


class BCEBlurWithLogitsLoss(paddle.nn.Layer):

    def __init__(self, alpha=0.05):
        super(BCEBlurWithLogitsLoss, self).__init__()
        self.loss_fcn = paddle.nn.BCEWithLogitsLoss(reduction='none')
        self.alpha = alpha

    def forward(self, pred, true):
        loss = self.loss_fcn(pred, true)
        pred = paddle.nn.functional.sigmoid(x=pred)
        dx = pred - true
        alpha_factor = 1 - paddle.exp(x=(dx - 1) / (self.alpha + 0.0001))
        loss *= alpha_factor
        return loss.mean()


class FocalLoss(paddle.nn.Layer):

    def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
        super(FocalLoss, self).__init__()
        self.loss_fcn = loss_fcn
        self.gamma = gamma
        self.alpha = alpha
        self.reduction = loss_fcn.reduction
        self.loss_fcn.reduction = 'none'

    def forward(self, pred, true):
        loss = self.loss_fcn(pred, true)
        pred_prob = paddle.nn.functional.sigmoid(x=pred)
        p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
        alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
        modulating_factor = (1.0 - p_t) ** self.gamma
        loss *= alpha_factor * modulating_factor
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:
            return loss


class QFocalLoss(paddle.nn.Layer):

    def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
        super(QFocalLoss, self).__init__()
        self.loss_fcn = loss_fcn
        self.gamma = gamma
        self.alpha = alpha
        self.reduction = loss_fcn.reduction
        self.loss_fcn.reduction = 'none'

    def forward(self, pred, true):
        loss = self.loss_fcn(pred, true)
        pred_prob = paddle.nn.functional.sigmoid(x=pred)
        alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
        modulating_factor = paddle.abs(x=true - pred_prob) ** self.gamma
        loss *= alpha_factor * modulating_factor
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:
            return loss


class ComputeLoss:

    def __init__(self, model, autobalance=False):
        self.sort_obj_iou = False
        device = next(model.parameters()).place
        h = model.hyp
        BCEcls = paddle.nn.BCEWithLogitsLoss(pos_weight=paddle.to_tensor(
            data=[h['cls_pw']], place=device))
        BCEobj = paddle.nn.BCEWithLogitsLoss(pos_weight=paddle.to_tensor(
            data=[h['obj_pw']], place=device))
        self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0))
        g = h['fl_gamma']
        if g > 0:
            BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
        det = model.module.model[-1] if is_parallel(model) else model.model[-1]
        self.balance = {(3): [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 
            0.06, 0.02])
        self.ssi = list(det.stride).index(16) if autobalance else 0
        self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = (
            BCEcls, BCEobj, 1.0, h, autobalance)
        for k in ('na', 'nc', 'nl', 'anchors'):
            setattr(self, k, getattr(det, k))

    def __call__(self, p, targets):
        device = targets.place
        lcls, lbox, lobj = paddle.zeros(shape=[1]), paddle.zeros(shape=[1]
            ), paddle.zeros(shape=[1])
        tcls, tbox, indices, anchors = self.build_targets(p, targets)
        for i, pi in enumerate(p):
            b, a, gj, gi = indices[i]
            tobj = paddle.zeros_like(x=pi[..., 0])
            n = tuple(b.shape)[0]
            if n:
                ps = pi[b, a, gj, gi]
                pxy = ps[:, :2].sigmoid() * 2.0 - 0.5
                pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
                pbox = paddle.concat(x=(pxy, pwh), axis=1)
                iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True)
                lbox += (1.0 - iou).mean()
                score_iou = iou.detach().clip(min=0).astype(tobj.dtype)
                if self.sort_obj_iou:
                    sort_id = paddle.argsort(x=score_iou)
                    b, a, gj, gi, score_iou = b[sort_id], a[sort_id], gj[
                        sort_id], gi[sort_id], score_iou[sort_id]
                tobj[b, a, gj, gi] = 1.0 - self.gr + self.gr * score_iou
                if self.nc > 1:
                    t = paddle.full_like(x=ps[:, 5:], fill_value=self.cn)
                    t[range(n), tcls[i]] = self.cp
                    lcls += self.BCEcls(ps[:, 5:], t)
            obji = self.BCEobj(pi[..., 4], tobj)
            lobj += obji * self.balance[i]
            if self.autobalance:
                self.balance[i] = self.balance[i
                    ] * 0.9999 + 0.0001 / obji.detach().item()
        if self.autobalance:
            self.balance = [(x / self.balance[self.ssi]) for x in self.balance]
        lbox *= self.hyp['box']
        lobj *= self.hyp['obj']
        lcls *= self.hyp['cls']
        bs = tuple(tobj.shape)[0]
        return (lbox + lobj + lcls) * bs, paddle.concat(x=(lbox, lobj, lcls)
            ).detach()

    def build_targets(self, p, targets):
        na, nt = self.na, tuple(targets.shape)[0]
        tcls, tbox, indices, anch = [], [], [], []
        gain = paddle.ones(shape=[7])
        ai = paddle.arange(end=na).astype(dtype='float32').view(na, 1).tile(
            repeat_times=[1, nt])
        targets = paddle.concat(x=(targets.tile(repeat_times=[na, 1, 1]),
            ai[:, :, None]), axis=2)
        g = 0.5
        off = paddle.to_tensor(data=[[0, 0], [1, 0], [0, 1], [-1, 0], [0, -
            1]], place=targets.place).astype(dtype='float32') * g
        for i in range(self.nl):
            anchors = self.anchors[i]
            gain[2:6] = paddle.to_tensor(data=tuple(p[i].shape))[[3, 2, 3, 2]]
            t = targets * gain
            if nt:
                r = t[:, :, 4:6] / anchors[:, None]
                j = paddle_aux.max(r, 1.0 / r).max(2)[0] < self.hyp['anchor_t']
                t = t[j]
                gxy = t[:, 2:4]
                gxi = gain[[2, 3]] - gxy
                j, k = ((gxy % 1.0 < g) & (gxy > 1.0)).T
                l, m = ((gxi % 1.0 < g) & (gxi > 1.0)).T
                j = paddle.stack(x=(paddle.ones_like(x=j), j, k, l, m))
                t = t.tile(repeat_times=(5, 1, 1))[j]
                offsets = (paddle.zeros_like(x=gxy)[None] + off[:, None])[j]
            else:
                t = targets[0]
                offsets = 0
            b, c = t[:, :2].astype(dtype='int64').T
            gxy = t[:, 2:4]
            gwh = t[:, 4:6]
            gij = (gxy - offsets).astype(dtype='int64')
            gi, gj = gij.T
            a = t[:, 6].astype(dtype='int64')
            indices.append((b, a, gj.clip_(min=0, max=gain[3] - 1), gi.
                clip_(min=0, max=gain[2] - 1)))
            tbox.append(paddle.concat(x=(gxy - gij, gwh), axis=1))
            anch.append(anchors[a])
            tcls.append(c)
        return tcls, tbox, indices, anch
